"""示例文件 - 演示如何使用重构后的代理系统。"""

import asyncio
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START
from langgraph.config import get_stream_writer
from langgraph.func import entrypoint, task
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
from agents.graphs import (
    ReactConfiguration,
    SqlConfiguration,
    react_graph,
    sql_graph,
)  # 移除了document_graph的导入

from utils import logger
from agents.graphs import create_react_sql_agent_async

# 导入统一的打印工具
from utils.pretty_printer import (
    printer,
    print_message,
    print_separator,
    print_header,
    print_step_info,
    print_error,
    print_success,
    print_warning,
)


def print_messages_colored(chunk):
    """打印消息块中的所有消息，使用统一的打印工具。"""
    if isinstance(chunk, dict) and "messages" in chunk:
        print_separator("消息更新")
        for msg in chunk["messages"]:
            print_message(msg)
        print_separator()
    else:
        print_message(str(chunk), "system")


def test_sql_agent():
    """测试 SQL 代理。"""
    print_header("测试 SQL 代理")

    # 创建消息
    messages = [HumanMessage(content="列出所有表格,并获取表格的详细信息。")]

    # 调用代理
    for chunk, metadata in sql_graph.invoke(
        {"messages": messages}, stream_mode="messages"
    ):
        print_message(f"Chunk: {chunk}", "system")
        print_message(f"Metadata: {metadata}", "system")
        # print_messages_colored(chunk)
        # print("\n")


import re


def format_response(response):
    """对输入的文本进行段落分隔、添加适当的换行符，以及在代码块中增加标记，以便生成更有可读性的输出。

    Args:
        response: 输入的文本。

    Returns:
        具有清晰段落分隔的文本。
    """
    # 使用正则表达式 \n{2, }将输入的response按照两个或更多的连续换行符进行分割。这样可以将文本分割成多个段落，每个段落由连续的非空行组成
    paragraphs = re.split(r"\n{2,}", response)
    # 空列表，用于存储格式化后的段落
    formatted_paragraphs = []
    # 遍历每个段落进行处理
    for para in paragraphs:
        # 检查段落中是否包含代码块标记
        if "```" in para:
            # 将段落按照```分割成多个部分，代码块和普通文本交替出现
            parts = para.split("```")
            for i, part in enumerate(parts):
                # 检查当前部分的索引是否为奇数，奇数部分代表代码块
                if i % 2 == 1:  # 这是代码块
                    # 将代码块部分用换行符和```包围，并去除多余的空白字符
                    parts[i] = f"\n```\n{part.strip()}\n```\n"
            # 将分割后的部分重新组合成一个字符串
            para = "".join(parts)
        else:
            # 否则，将句子中的句点后面的空格替换为换行符，以便句子之间有明确的分隔
            para = para.replace(". ", ".\n")
        # 将格式化后的段落添加到formatted_paragraphs列表
        # strip()方法用于移除字符串开头和结尾的空白字符（包括空格、制表符 \t、换行符 \n等）
        formatted_paragraphs.append(para.strip())
    # 将所有格式化后的段落用两个换行符连接起来，以形成一个具有清晰段落分隔的文本
    return "\n\n".join(formatted_paragraphs)


async def test_react_agent():
    """测试 React 代理。"""
    print_header("测试 React 代理")

    # 配置 React 代理
    config = ReactConfiguration(
        model_name="gpt-4o-mini",
        temperature=0.7,
        max_tokens=1000,
    )

    # 创建消息
    messages = [HumanMessage(content="今天天气怎么样？")]

    # 调用代理
    result = await react_graph.ainvoke(
        {"messages": messages},
        config={"configurable": config.model_dump()},
    )

    print_message(f"React代理响应: {result['messages'][-1].content}", "ai")


from IPython.display import display, Image
import os


# 将构建的graph可视化保存为 PNG 文件
def save_graph_visualization(
    graph: StateGraph, filename: str = "sql_graph.png"
) -> None:
    try:
        os.makedirs("output", exist_ok=True)
        file_path = os.path.join("output", filename)
        with open(file_path, "wb") as f:
            f.write(graph.get_graph().draw_mermaid_png())
        print_success(f"Graph visualization saved as {file_path}")
    except IOError as e:
        print_warning(f"Failed to save graph visualization: {str(e)}")


async def test_sql_agent_async_non_streaming():
    """测试 SQL 代理 - 异步非流式输出。"""
    print_header("测试 SQL 代理 (异步非流式)")

    try:
        # 导入异步图创建函数
        from agents.graphs import create_sql_graph_async

        # 创建异步SQL图
        async_sql_graph = await create_sql_graph_async()

        # 创建消息
        messages = [HumanMessage(content="列出所有表格,并获取表格的详细信息。")]

        # 使用异步接口调用代理 - 非流式输出
        config = {"configurable": {"thread_id": "test_thread_1"}}
        result = await async_sql_graph.ainvoke({"messages": messages}, config=config)

        print_separator("非流式输出结果")
        for message in result["messages"]:
            print_message(message)

        print_success("测试完成")

    except Exception as e:
        print_error(f"测试失败: {str(e)}")
        logger.error(f"Async non-streaming test failed: {e}")


async def test_sql_agent_async_streaming():
    """测试 SQL 代理 - 异步流式输出。"""
    print_header("测试 SQL 代理 (异步流式)")

    try:
        # 导入异步图创建函数
        from agents.graphs import create_sql_graph_async

        # 创建异步SQL图
        async_sql_graph = await create_sql_graph_async()

        # 创建消息
        messages = [HumanMessage(content="列出所有表格,并获取表格的详细信息。")]

        # 使用异步接口调用代理 - 流式输出
        config = {"configurable": {"thread_id": "test_thread_2"}}

        print_separator("流式输出结果")
        async for chunk, metadata in async_sql_graph.astream(
            {"messages": messages}, config=config, stream_mode="messages"
        ):
            print_message(f"Chunk: {chunk}", "system")
            print_message(f"Metadata: {metadata}", "system")
            print_message(chunk)
            print_separator()

        print_success("测试完成")

    except Exception as e:
        print_error(f"测试失败: {str(e)}")
        logger.error(f"Async streaming test failed: {e}")


def print_step_details(step, step_num=None):
    """打印步骤详细信息的辅助函数"""
    if step_num:
        print_step_info(step_num, "处理步骤")

    if isinstance(step, dict):
        if "messages" in step and step["messages"]:
            print_separator("消息内容")
            for i, message in enumerate(step["messages"]):
                print_message(f"消息 {i+1}:", "system")
                print_message(message)
        else:
            print_message(f"步骤数据: {step}", "system")
    else:
        print_message(f"步骤内容: {step}", "system")


async def main():
    # save_graph_visualization(sql_graph)
    """主函数。"""
    print_header("LangGraph 代理系统测试")

    # 测试同步版本（可能会有错误）
    # test_sql_agent()

    # 测试异步非流式输出
    # await test_sql_agent_async_non_streaming()

    try:
        print_step_info(1, "创建 React SQL 代理")
        agent = await create_react_sql_agent_async()
        print_success("代理创建成功")

        question = (
            "年度NCM或fix time排在top10的物料有哪些? 相关数据源表是ModelCapaSsme,基础表"
        )
        print_message(f"用户问题: {question}", "human")

        config = {
            "configurable": {"thread_id": "thread_id", "user_id": "user_id"},
            "recursion_limit": 25,
        }

        print_separator("代理执行过程")
        step_count = 0

        async for chunk, metadata in agent.astream(
            {"messages": [{"role": "user", "content": question}]},
            stream_mode="messages",
            config=config,
        ):
            step_count += 1
            print_step_info(step_count, f"处理消息块")
            print(chunk)
            # print_message(f"Chunk: {chunk}", "system")
            # print_message(f"Metadata: {metadata}", "system")
            # if hasattr(chunk, 'content') and chunk.content:
            #     print_message(chunk)
            # print_separator()

        print_success(f"代理执行完成，共处理 {step_count} 个消息块")

    except Exception as e:
        print_error(f"执行失败: {str(e)}")
        logger.error(f"Main execution failed: {e}")

    # 测试异步流式输出
    # await test_sql_agent_async_streaming()

    # await test_react_agent()

    # 注释掉document_graph相关的测试代码
    # result = document_graph.invoke(
    #     {"messages": [HumanMessage(content="分析这个文档")]},
    #     config={"configurable": {}},
    # )
    # print(f"Document代理响应: {result['messages'][-1].content}")


if __name__ == "__main__":
    asyncio.run(main())  # 使用asyncio.run运行异步主函数
